# langchain_agent_demo.py：适配 langchain==0.3.27 的 Agent 代码
from langchain.tools import tool
# 1. 替换 ChatOpenAI 导入（从 langchain-openai）
from langchain_openai import ChatOpenAI
from langchain.agents import ZeroShotAgent, AgentExecutor
from langchain.prompts import PromptTemplate
from dotenv import load_dotenv
import os

# 加载环境变量
load_dotenv()

# 工具定义部分不变（WeatherTool、CalculatorTool、RAGTool）
@tool(description="用于查询指定城市的当前天气。参数：city（必填，城市名，示例：北京、上海），返回值：温度和天气状况。用户问城市天气必须调用此工具，不可编造数据。")
def WeatherTool(city: str) -> str:
    weather_data = {
        "北京": "温度18-25℃，晴，微风",
        "上海": "温度20-28℃，多云，东南风3级",
        "广州": "温度25-32℃，阴，有小雨"
    }
    return weather_data.get(city, f"暂未获取到{city}的天气数据")

@tool(description="用于计算加减乘除运算。参数：expression（必填，数学表达式，示例：2+3*4、10-5，无空格），返回值：计算结果。仅处理数字和+、-、*、/。")
def CalculatorTool(expression: str) -> str:
    try:
        result = eval(expression)
        return f"计算结果：{expression} = {result}"
    except Exception as e:
        return f"计算失败：{str(e)}，请检查表达式（示例：2+3，不要有空格）"

@tool(description="用于查询公司内部知识库（如年假政策、产品价格）。参数：question（必填，查询问题，示例：公司年假政策、产品A价格），返回值：知识库答案。公司相关问题必须调用此工具，不可编造。")
def RAGTool(question: str) -> str:
    rag_data = {
        "公司年假政策": "员工入职满1年5天年假，满3年10天，满5年15天。",
        "产品A的价格": "产品A基础版999元/年，高级版1999元/年，企业版需定制。",
        "加班补贴标准": "工作日20元/小时，周末30元/小时，法定节假日50元/小时。"
    }
    for key, value in rag_data.items():
        if key in question:
            return value
    return f"知识库中未找到「{question}」的信息，请联系相关部门。"

# 整理工具列表
tools = [WeatherTool, CalculatorTool, RAGTool]


# 2. 初始化 LLM（使用新的导入方式）
llm = ChatOpenAI(
    api_key=os.getenv("DASHSCOPE_API_KEY"),
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
    model="qwen-turbo",
)

# 3. 修复提示词模板（添加工具名称变量 {tool_names}）
prompt_template = """
你是帮助用户解决问题的智能Agent，可使用以下工具：
{tools}
可用工具名称：{tool_names}  # 新增工具名称列表

使用规则：
1. 先理解问题：无需工具则直接回答，需要工具则选对应工具（根据工具描述匹配）；
2. 调用工具必须用格式：```json{{"name":"工具名","parameters":{{"参数名":"参数值"}}}}```（参数需齐全）；
3. 拿到工具结果后，判断是否需要继续调用工具，或整理结果回答。

用户问题：{input}
思考过程：{agent_scratchpad}
"""

# 创建提示词
prompt = PromptTemplate(
    template=prompt_template,
    input_variables=["tools", "tool_names", "input", "agent_scratchpad"]  # 新增 tool_names 变量
)

# 4. 修复 Agent 初始化方式（使用 from_llm_and_tools 工厂方法）
agent = ZeroShotAgent.from_llm_and_tools(
    llm=llm,
    tools=tools,
    prompt=prompt,
    verbose=True
)



# 创建 Agent 执行器
agent_executor = AgentExecutor(
    agent=agent,
    tools=tools,
    verbose=True,
    handle_parsing_errors=True,
    max_iterations=5,

)

# 测试场景
if __name__ == "__main__":

    # # 手动填充所有参数（模拟Agent运行时的真实参数）
    # full_prompt = prompt.format(
    #     tools=tools,  # 工具列表
    #     tool_names=", ".join([tool.name for tool in tools]),  # 工具名称列表（和框架自动填充一致）
    #     chat_history="",  # 历史对话（首次调用为空）
    #     input="查询苏州天气",  # 示例用户问题
    #     agent_scratchpad=""  # 初始思考过程为空
    # )
    #
    # # 打印完整提示词
    # print("=" * 50)
    # print("完整提示词内容：")
    # print(full_prompt)
    # print("=" * 50)
    # print("=" * 50)
    # print("场景1：无需工具 → 什么是人工智能？")
    # result1 = agent_executor.invoke({"input": "什么是人工智能？用简单的话解释。"})  # 传入字典格式的input
    # print("答案：", result1["output"])  # 从返回结果的"output"字段提取答案
    #
    # # 场景2：单个工具（查北京天气）
    # print("\n" + "=" * 50)
    # print("场景2：单个工具 → 查询北京天气")
    # result2 = agent_executor.invoke({"input": "查询北京今天的天气。"})
    # print("答案：", result2["output"])

    print("\n" + "="*50)
    print("场景3：多个工具 → 上海天气+穿搭+预算")
    result3 = agent_executor.invoke({"input": "查上海今天的天气，推荐适合的穿搭，再计算这套穿搭的预算（T恤150元，牛仔裤200元，外套300元）。"})
    print("答案：", result3)

    # print("\n" + "="*50)
    # print("场景4：RAG工具 → 公司年假政策")
    # result4 = agent_executor.invoke({"input": "公司员工入职满3年有多少天年假？"})
    # print("答案：", result4)